import cv2
import time
import numpy as np
from keras.layers import Input
from keras.models import Model
from PIL import Image,ImageDraw,ImageFont
from yolov3.yolov3_body import body

def load_coconames(file_name):
    with open(file_name) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]

    return class_names

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def get_ypre(ypre,input_size,anchors):
    rate = input_size/int(ypre.shape[1])
    ypre = np.reshape(ypre,[ypre.shape[-3], ypre.shape[-2], 3, ypre.shape[-1]//3])

    mesh_x = np.reshape(np.tile(np.arange(ypre.shape[0]), [ypre.shape[1]]), ( ypre.shape[0], ypre.shape[1], 1, 1))
    mesh_y = np.transpose(mesh_x, (1,0,2,3))
    mesh_xy = np.tile(np.concatenate([mesh_x,mesh_y],-1), [1, 1,ypre.shape[2], 1])

    anchor_gird = np.reshape(np.array(anchors),(1,1,3,2))

    pre_xy = (sigmoid(ypre[...,:2])+mesh_xy)*rate
    pre_wh = np.exp(ypre[...,2:4])*anchor_gird
    pre_con = sigmoid(ypre[...,4:5])
    pre_class = sigmoid(ypre[...,5:])

    pre_wh_half = pre_wh/2
    pre_mins = pre_xy-pre_wh_half
    pre_maxes = pre_xy+pre_wh_half

    return np.concatenate([pre_mins,pre_maxes,pre_con,pre_class],axis=-1)

def get_IOU(box1,box2):
    intersect_mins  =np.maximum(box1[:2],  box2[:2])
    intersect_maxes =np.minimum(box1[2:4],  box2[2:4])

    intersect_wh    = np.maximum(intersect_maxes - intersect_mins, 0.)
    intersect_areas = intersect_wh[..., 0] * intersect_wh[..., 1]

    box1_areas = (box1[2]-box1[0]) * (box1[3]-box1[1])
    box2_areas = (box2[2]-box2[0]) * (box2[3]-box2[1])

    union_areas = box1_areas + box2_areas - intersect_areas
    IOU  = intersect_areas/union_areas

    return IOU

def nms_boxes(box,iou):
    boxes = sorted(box.tolist(),key=(lambda x:x[5]),reverse=True)
    new_boxes = []
    l = 0
    while len(boxes)>0:
        new_boxes.append(boxes[0])
        del boxes[0]
        i = 0
        for _ in range(len(boxes)):
            IOU = get_IOU(np.array(new_boxes[l][:4]), np.array(boxes[i][:4]))
            if IOU>iou:
                del boxes[i]
                i -= 1
            i += 1
        l += 1
    return new_boxes

def draw_result(boxes,img,class_names,out):
    draw = ImageDraw.Draw(img)
    for box in boxes:
        xmin = box[0] * img.size[0] / 416
        ymin = box[1] * img.size[1] / 416
        xmax = box[2] * img.size[0] / 416
        ymax = box[3] * img.size[1] / 416

        draw.rectangle((xmin, ymin, xmax, ymax),fill=None, outline=(237,106,62), width=1)

        font = ImageFont.truetype("consola.ttf", 40, encoding="unic")
        label = str(class_names[int(box[6])])+' '+str(format(box[5]*100,'.2f'))+'%'
        draw.text((xmin, ymin), label, fill=(237, 106, 62), font=font)

    img = cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

    cv2.imshow("camera",img)
    cv2.waitKey(1)

    out.write(img)

def main():
    iou = 0.5
    score = 0.3
    anchors = [[[10, 13], [16, 30], [33, 23]], [[30, 61], [62, 45], [59, 119]], [[116, 90], [156, 198], [373, 326]]]
    class_names = load_coconames('D:/py/yolov3/coco_classes.txt')

    model_input = Input(shape=(None,None,3))
    model_output = body(model_input,3,len(class_names))
    model = Model(model_input,model_output)
    print("已构建模型")
    model.load_weights('D:/py/yolov3/yolo_weights.h5')
    print("已加载权重")

    cap = cv2.VideoCapture('D:/py/yolov3/in.mp4')
    print("已读取视频")
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter('out.mp4', fourcc, 30, (1920, 1080), True)
    while True:
        start_time = time.time()

        ret,frame = cap.read()
        if not ret:
            break
        img = Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB))
        image = img.resize(size=(416,416))
        image_data = np.array(image,dtype='float32')
        image_data /= 255.
        image_data = np.expand_dims(image_data,0)

        ypres = model.predict(image_data)
        boxes = []
        l = 2
        for ypre in ypres:
            new_ypre = get_ypre(ypre,416,anchors[l])
            for m in range(int(ypre.shape[2])):
                for n in range(int(ypre.shape[1])):
                    for i in range(3):
                        object_prob = new_ypre[m,n,i,4]
                        classes_prob = new_ypre[m,n,i,5:]
                        classes = classes_prob*object_prob
                        classes = classes.tolist()
                        index = classes.index(max(classes))

                        new_ypre[m,n,i,5:7] = (max(classes),index)

                        if object_prob>score:
                            boxes.append(new_ypre[m,n,i,:7])

            l -= 1

        boxes = nms_boxes(np.array(boxes),iou)

        draw_result(boxes,img,class_names,out)

        end_time = time.time()
        print('单帧处理时间: {:.3f}s'.format(end_time - start_time))


if __name__ == '__main__':
    main()